% 2-D diffusion solver, fixed source, one group
% multi material; tests the basics and with reflective
% bc's can be used to test against working 1-D code
clear; clc;
% example 10 by 10 square with four different media
N = 9;
% basic material data
data = [ 0.6500    0.140    0.175    0.000        
         0.2000    0.020    0.000    0.000 
         0.6000    0.450    0.170    0.000];
% put the material in their physical location
mat = 2*ones(10);
% mat =  [ 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2
%          2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2
%          2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2
%          2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2
%          2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2
%          2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2
%          3 3 3 3 3 3 2 2 2 2 2 2 2 2 2 2 2 2 2 2
%          3 3 3 3 3 3 2 2 2 2 2 2 2 2 2 2 2 2 2 2
%          2 2 2 2 3 3 2 2 2 2 2 2 2 2 2 2 2 2 2 2
%          2 2 2 2 3 3 2 2 2 2 2 2 2 2 2 2 2 2 2 2
%          2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 
%          2 2 2 2 3 3 3 2 2 2 2 2 2 2 2 2 2 2 2 2
%          3 3 3 3 3 3 2 2 2 2 2 2 2 2 2 2 2 2 2 2
%          3 3 3 3 3 3 2 2 2 2 2 2 2 2 2 2 2 2 2 2
%          2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2
%          2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2
%          2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2
%          2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2
%          2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2
%          2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2];
% create "material" matrices, probably not necessary
for i = 1:N+1
    for j = 1:N+1
        difc(i,j) = data(mat(i,j),1);
        siga(i,j) = data(mat(i,j),2);
        nusi(i,j) = data(mat(i,j),3);
        widx(i,j) = 1.0;
        widy(i,j) = 1.0;
    end
end
% define source
src = ones(10);
% src =  [ 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
%          0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
%          0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
%          0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
%          0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
%          0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
%          0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
%          0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
%          1 1 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
%          1 1 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
%          1 1 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
%          1 1 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
%          0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
%          0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
%          0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
%          0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
%          0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
%          0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
%          0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
%          0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 ];
% define the "A" matrix
% a given element (i,j) connects to (i-1,j), (i+1,j), (i,j-1), (i,j+1)
% each (i,j) is assigned to a unique k
A = zeros(N*N); S = zeros(N,1);
% deal with the internal ones
for i = 2:N-1
    for j = 2:N-1
        k = i+(j-1)*( N ); % 10 being the number of x ints
        % now make the "A" matrix
        A(k, i-1+(j-1)*( N )) = ...
            -0.5*(difc(i,j)*widy(i,j)+difc(i,j+1)*widy(i,j+1))/widx(i,j);
        A(k, i+1+(j-1)*( N )) = ...
            -0.5*(difc(i+1,j)*widy(i+1,j)+difc(i+1,j+1)*widy(i+1,j+1))/widx(i+1,j);
        A(k, i+(j-1-1)*( N )) = ...
            -0.5*(difc(i,j)*widx(i,j)+difc(i+1,j)*widx(i+1,j))/widy(i,j);
        A(k, i+(j-1+1)*( N )) = ...
            -0.5*(difc(i,j+1)*widx(i,j+1)+difc(i+1,j+1)*widx(i+1,j+1))/widy(i,j+1);
        A(k,k) = -( A(k, i-1+(j-1)*( N ))+A(k, i+1+(j-1)*( N ))+ ...
            A(k, i+(j-1-1)*( N ))+A(k, i+(j-1+1)*( N )) )+ ...
            0.25*( widx(i,j)*widy(i,j)*siga(i,j) + widx(i+1,j)*widy(i+1,j)*siga(i+1,j) + ...
            widx(i,j+1)*widy(i,j+1)*siga(i,j+1) + widx(i+1,j+1)*widy(i+1,j)*siga(i+1,j+1) );
        S(k)   = 0.25*( widx(i,j)*widy(i,j)*src(i,j) + widx(i+1,j)*widy(i+1,j)*src(i+1,j) + ...
            widx(i,j+1)*widy(i,j+1)*src(i,j+1) + widx(i+1,j+1)*widy(i+1,j)*src(i+1,j+1) );
    end
end

% deal with the edges, not corners
for j = 2:N-1
    i = 1;
    k = i+(j-1)*( N );
    A(k, i+1+(j-1)*( N )) = ... 
        -0.5*(difc(i+1,j)*widy(i+1,j)+difc(i+1,j+1)*widy(i+1,j+1))/widx(i+1,j);
    A(k, i+(j-1-1)*( N )) = ... 
        -0.5*(difc(i,j)*widx(i,j)+difc(i+1,j)*widx(i+1,j))/widy(i,j);
    A(k, i+(j-1+1)*( N )) = ...
        -0.5*(difc(i,j+1)*widx(i,j+1)+difc(i+1,j+1)*widx(i+1,j+1))/widy(i,j+1);
    A(k,k) = -( A(k, i+1+(j-1)*( N ))+ ...
        A(k, i+(j-1-1)*( N ))+A(k, i+(j-1+1)*( N )) )+ ...
        0.25*(  widx(i+1,j)*widy(i+1,j)*siga(i+1,j) + ...
        widx(i,j+1)*widy(i,j+1)*siga(i,j+1) + widx(i+1,j+1)*widy(i+1,j)*siga(i+1,j+1) );    
    S(k) = 0.25*( widx(i,j)*widy(i,j)*src(i,j) + widx(i+1,j)*widy(i+1,j)*src(i+1,j) + ...
        widx(i,j+1)*widy(i,j+1)*src(i,j+1) + widx(i+1,j+1)*widy(i+1,j)*src(i+1,j+1) );   
    i = N;
    k = i+(j-1)*( N );
    A(k, i-1+(j-1)*( N )) = ...
        -0.5*(difc(i,j)*widy(i,j)+difc(i,j+1)*widy(i,j+1))/widx(i,j);
    A(k, i+(j-1-1)*( N )) = ...
        -0.5*(difc(i,j)*widx(i,j)+difc(i+1,j)*widx(i+1,j))/widy(i,j);
    A(k, i+(j-1+1)*( N )) = ...
        -0.5*(difc(i,j+1)*widx(i,j+1)+difc(i+1,j+1)*widx(i+1,j+1))/widy(i,j+1);       
    A(k,k) = -( A(k, i-1+(j-1)*( N ))+ ...
        A(k, i+(j-1-1)*( N ))+A(k, i+(j-1+1)*( N )) )+ ...
        0.25*( widx(i,j)*widy(i,j)*siga(i,j) + widx(i+1,j)*widy(i+1,j)*siga(i+1,j) + ...
        widx(i,j+1)*widy(i,j+1)*siga(i,j+1) + widx(i+1,j+1)*widy(i+1,j)*siga(i+1,j+1) );    
    S(k) = 0.25*( widx(i,j)*widy(i,j)*src(i,j) + widx(i+1,j)*widy(i+1,j)*src(i+1,j) + ...
        widx(i,j+1)*widy(i,j+1)*src(i,j+1) + widx(i+1,j+1)*widy(i+1,j)*src(i+1,j+1) );
end

for i = 2:N-1
    j = 1;
    k = i+(j-1)*( N );
    A(k, i-1+(j-1)*( N )) = ...
        -0.5*(difc(i,j)*widy(i,j)+difc(i,j+1)*widy(i,j+1))/widx(i,j);
    A(k, i+1+(j-1)*( N )) = ...
        -0.5*(difc(i+1,j)*widy(i+1,j)+difc(i+1,j+1)*widy(i+1,j+1))/widx(i+1,j);
    A(k, i+(j-1+1)*( N )) = ...
        -0.5*(difc(i,j+1)*widx(i,j+1)+difc(i+1,j+1)*widx(i+1,j+1))/widy(i,j+1);
    A(k,k) = -( A(k, i-1+(j-1)*( N ))+A(k, i+1+(j-1)*( N ))+ ...
        A(k, i+(j-1+1)*( N )) )+ ...
        0.25*( widx(i,j)*widy(i,j)*siga(i,j) + widx(i+1,j)*widy(i+1,j)*siga(i+1,j) + ...
        widx(i,j+1)*widy(i,j+1)*siga(i,j+1) + widx(i+1,j+1)*widy(i+1,j)*siga(i+1,j+1) );
    S(k)   = 0.25*( widx(i,j)*widy(i,j)*src(i,j) + widx(i+1,j)*widy(i+1,j)*src(i+1,j) + ...
        widx(i,j+1)*widy(i,j+1)*src(i,j+1) + widx(i+1,j+1)*widy(i+1,j)*src(i+1,j+1) );
    j = N;
    k = i+(j-1)*( N );
    A(k, i-1+(j-1)*( N )) = ...
        -0.5*(difc(i,j)*widy(i,j)+difc(i,j+1)*widy(i,j+1))/widx(i,j);
    A(k, i+1+(j-1)*( N )) = ...
        -0.5*(difc(i+1,j)*widy(i+1,j)+difc(i+1,j+1)*widy(i+1,j+1))/widx(i+1,j);
    A(k, i+(j-1-1)*( N )) = ...
        -0.5*(difc(i,j)*widx(i,j)+difc(i+1,j)*widx(i+1,j))/widy(i,j); 
    A(k,k) = -( A(k, i-1+(j-1)*( N ))+A(k, i+1+(j-1)*( N ))+ ...
        A(k, i+(j-1-1)*( N )) )+ ...
        0.25*( widx(i,j)*widy(i,j)*siga(i,j) + widx(i+1,j)*widy(i+1,j)*siga(i+1,j) + ...
        widx(i,j+1)*widy(i,j+1)*siga(i,j+1) + widx(i+1,j+1)*widy(i+1,j)*siga(i+1,j+1) );
    S(k)   = 0.25*( widx(i,j)*widy(i,j)*src(i,j) + widx(i+1,j)*widy(i+1,j)*src(i+1,j) + ...
        widx(i,j+1)*widy(i,j+1)*src(i,j+1) + widx(i+1,j+1)*widy(i+1,j)*src(i+1,j+1) );
end

% finally, the corners
i=1; j=1; k = i+(j-1)*( N );
A(k, i+1+(j-1)*( N )) = ... % R
    -0.5*(difc(i+1,j)*widy(i+1,j)+difc(i+1,j+1)*widy(i+1,j+1))/widx(i+1,j);
A(k, i+(j-1+1)*( N )) = ... % T
    -0.5*(difc(i,j+1)*widx(i,j+1)+difc(i+1,j+1)*widx(i+1,j+1))/widy(i,j+1);
A(k,k) = -( A(k, i+1+(j-1)*( N ))+ A(k, i+(j-1+1)*( N )) )+ ...
    0.25*( widx(i,j)*widy(i,j)*siga(i,j) + widx(i+1,j)*widy(i+1,j)*siga(i+1,j) + ...
    widx(i,j+1)*widy(i,j+1)*siga(i,j+1) + widx(i+1,j+1)*widy(i+1,j)*siga(i+1,j+1) );
S(k)   = 0.25*( widx(i,j)*widy(i,j)*src(i,j) + widx(i+1,j)*widy(i+1,j)*src(i+1,j) + ...
    widx(i,j+1)*widy(i,j+1)*src(i,j+1) + widx(i+1,j+1)*widy(i+1,j)*src(i+1,j+1) );
%
i=1; j=N; k = i+(j-1)*( N );
A(k, i+1+(j-1)*( N )) = ... % R
    -0.5*(difc(i+1,j)*widy(i+1,j)+difc(i+1,j+1)*widy(i+1,j+1))/widx(i+1,j);
A(k, i+(j-1-1)*( N )) = ... % B
    -0.5*(difc(i,j)*widx(i,j)+difc(i+1,j)*widx(i+1,j))/widy(i,j);
A(k,k) = -( A(k, i+1+(j-1)*( N ))+A(k, i+(j-1-1)*( N )) )+ ...
    0.25*( widx(i,j)*widy(i,j)*siga(i,j) + widx(i+1,j)*widy(i+1,j)*siga(i+1,j) + ...
    widx(i,j+1)*widy(i,j+1)*siga(i,j+1) + widx(i+1,j+1)*widy(i+1,j)*siga(i+1,j+1) );
S(k)   = 0.25*( widx(i,j)*widy(i,j)*src(i,j) + widx(i+1,j)*widy(i+1,j)*src(i+1,j) + ...
    widx(i,j+1)*widy(i,j+1)*src(i,j+1) + widx(i+1,j+1)*widy(i+1,j)*src(i+1,j+1) );
%
i=N; j=1; k = i+(j-1)*( N );
A(k, i-1+(j-1)*( N )) = ... % L
    -0.5*(difc(i,j)*widy(i,j)+difc(i,j+1)*widy(i,j+1))/widx(i,j);
A(k, i+(j-1+1)*( N )) = ... % T
    -0.5*(difc(i,j+1)*widx(i,j+1)+difc(i+1,j+1)*widx(i+1,j+1))/widy(i,j+1);
A(k,k) = -( A(k, i-1+(j-1)*( N ))+A(k, i+(j-1+1)*( N )) )+ ...
    0.25*( widx(i,j)*widy(i,j)*siga(i,j) + widx(i+1,j)*widy(i+1,j)*siga(i+1,j) + ...
    widx(i,j+1)*widy(i,j+1)*siga(i,j+1) + widx(i+1,j+1)*widy(i+1,j)*siga(i+1,j+1) );
S(k)   = 0.25*( widx(i,j)*widy(i,j)*src(i,j) + widx(i+1,j)*widy(i+1,j)*src(i+1,j) + ...
    widx(i,j+1)*widy(i,j+1)*src(i,j+1) + widx(i+1,j+1)*widy(i+1,j)*src(i+1,j+1) );
%
i=N; j=N; k = i+(j-1)*( N );
A(k, i-1+(j-1)*( N )) = ... % L
    -0.5*(difc(i,j)*widy(i,j)+difc(i,j+1)*widy(i,j+1))/widx(i,j);
A(k, i+(j-1-1)*( N )) = ... % B
    -0.5*(difc(i,j)*widx(i,j)+difc(i+1,j)*widx(i+1,j))/widy(i,j);
A(k,k) = -( A(k, i-1+(j-1)*( N ))+A(k, i+(j-1-1)*( N )))+ ...
    0.25*( widx(i,j)*widy(i,j)*siga(i,j) + widx(i+1,j)*widy(i+1,j)*siga(i+1,j) + ...
    widx(i,j+1)*widy(i,j+1)*siga(i,j+1) + widx(i+1,j+1)*widy(i+1,j)*siga(i+1,j+1) );
S(k)   = 0.25*( widx(i,j)*widy(i,j)*src(i,j) + widx(i+1,j)*widy(i+1,j)*src(i+1,j) + ...
    widx(i,j+1)*widy(i,j+1)*src(i,j+1) + widx(i+1,j+1)*widy(i+1,j)*src(i+1,j+1) );

%% now solve
phi = A\S;
% generate indices
for i = 1:N
    for j = 1:N
        k = i+(j-1)*( N );
        indx(k,1)=i;
        indx(k,2)=j;
    end
end
% generate x and y vals
xint = zeros(N+2,1); yint=xint;
for i = 2:N+2
    xint(i)=xint(i-1)+widx(i-1,1);
    yint(i)=yint(i-1)+widy(i-1,1);
end
phimap = zeros(N+2);
for k = 1:N*N
    i = indx(k,1)+1;
    j = indx(k,2)+1;
    phimap(i,j)=phi(k);
end

surf(xint,yint,phimap')
xlabel('y')
ylabel('x')